import torch
import torch.nn as nn
import torch.nn.functional as F


class FreqEncoder_torch(nn.Module):
    def __init__(
        self,
        input_dim,
        max_freq_log2,
        N_freqs,
        log_sampling=True,
        include_input=True,
        periodic_fns=(torch.sin, torch.cos),
    ):
        super().__init__()

        self.input_dim = input_dim
        self.include_input = include_input
        self.periodic_fns = periodic_fns
        self.N_freqs = N_freqs
        self.output_dim = 0
        if self.include_input:
            self.output_dim += self.input_dim

        self.output_dim += self.input_dim * N_freqs * len(self.periodic_fns)

        if log_sampling:
            self.freq_bands = 2.0 ** torch.linspace(0.0, max_freq_log2, N_freqs)
        else:
            self.freq_bands = torch.linspace(2.0**0.0, 2.0**max_freq_log2, N_freqs)

        self.freq_bands = self.freq_bands.numpy().tolist()

    def forward(self, input, step_id=None):
        out = []
        if self.include_input:
            out.append(input)

        for i in range(len(self.freq_bands)):
            freq = self.freq_bands[i]
            for p_fn in self.periodic_fns:
                out.append(p_fn(input * freq))

        out = torch.cat(out, dim=-1)

        if step_id is not None:
            # use step_id to progressively open the frequency
            step_interval = 100
            soft_mask = [0] * self.N_freqs
            soft_mask[0], soft_mask[1] = 1, 1
            for i in range(2, self.N_freqs):
                if step_id < (i - 1) * step_interval:
                    break
                soft_mask[i] += min(
                    (step_id - (i - 1) * step_interval) / step_interval, 1
                )
            mask = torch.ones(self.output_dim)
            for i in range(self.N_freqs):
                idx = (
                    len(self.periodic_fns) * self.input_dim * i + self.input_dim
                    if self.include_input
                    else len(self.periodic_fns) * self.input_dim * i
                )
                mask[idx : idx + len(self.periodic_fns) * self.input_dim] *= soft_mask[
                    i
                ]
            mask = mask.to(out.device)
            out *= mask
        return out


def get_encoder(
    encoding,
    input_dim=3,
    multires=6,  # freq
    degree=4,  # SH
    num_levels=16,
    level_dim=2,
    base_resolution=16,
    log2_hashmap_size=19,
    desired_resolution=2048,  # hash/tiled grid
    align_corners=False,
    interpolation="linear",  # grid
    # interpolation="smoothstep",
    **kwargs
):
    if encoding == "None":
        return lambda x, **kwargs: x, input_dim

    elif encoding == "frequency_torch":
        encoder = FreqEncoder_torch(
            input_dim=input_dim,
            max_freq_log2=multires - 1,
            N_freqs=multires,
            log_sampling=True,
        )

    elif encoding == "frequency":
        from freqencoder import FreqEncoder

        encoder = FreqEncoder(input_dim=input_dim, degree=multires)

    elif encoding == "sh":
        from shencoder import SHEncoder

        encoder = SHEncoder(input_dim=input_dim, degree=degree)

    elif encoding == "hashgrid":
        from model.gridencoder import GridEncoder

        encoder = GridEncoder(
            input_dim=input_dim,
            num_levels=num_levels,
            level_dim=level_dim,
            base_resolution=base_resolution,
            log2_hashmap_size=log2_hashmap_size,
            desired_resolution=desired_resolution,
            gridtype="hash",
            align_corners=align_corners,
            interpolation=interpolation,
        )

    elif encoding == "tiledgrid":
        from gridencoder import GridEncoder

        encoder = GridEncoder(
            input_dim=input_dim,
            num_levels=num_levels,
            level_dim=level_dim,
            base_resolution=base_resolution,
            log2_hashmap_size=log2_hashmap_size,
            desired_resolution=desired_resolution,
            gridtype="tiled",
            align_corners=align_corners,
            interpolation=interpolation,
        )

    else:
        raise NotImplementedError(
            "Unknown encoding mode, choose from [None, frequency, sh, hashgrid, tiledgrid]"
        )

    return encoder, encoder.output_dim
